// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file 
\brief Implementation of Phenotype Utility Functions (Text File) 
*/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "copyfile.h"
#include "phefile.h"
#include "suffixes.h"
#include "utils.h"

#define MAX_LINE_LENGTH 200000
#define MAX_TOKEN_LENGTH 80

// Use array indexes to merge the contents of the PHE file with the PED file.
int phefile_read(struct pedfile *ped, struct selected_options *ops) {
	FILE *fin = NULL;
	char line[MAX_LINE_LENGTH];
	int i = 0, sampleIndex = 0;
	char msg[MAX_LINE_LENGTH];
	double missing_phenotype = 0;

	if (ped == NULL || ops == NULL) return 0;
	else if (ped->nSample == 0) return 0;
	// Ignorable ID columns from the PLINK PHE file format
	else if (ops->mpheno < 1 || strlen(ops->szPheFileName) == 0) return 0; 

	missing_phenotype = atof(ops->missing_phenotype);

	// Assume binary trait unless we find out otherwise.
	ops->bt = TRUE;
	ops->qt = FALSE;

	// Set all the sample to missing until phenotype is assigned.
	for (i = 0; i < ped->nSample; i++) {
		ped->sample[i]->missing = TRUE;
	}

	fin = fopen(ops->szPheFileName, "r");
	if (fin == NULL) return 0;

	// Echo message that file is loaded.
	if (ops->verbose) {
		sprintf(msg, "Reading PHE file: %s", ops->szPheFileName);
		print_log(msg);
	}

	// Read in a line, tokenise then get the bit we want.
	// If the token count does not agree with the index specified by ops->mpheno,
	// close and exit.
	while(fgets(line,MAX_LINE_LENGTH, fin)) {
		char token[MAX_TOKEN_LENGTH];
		char *p = NULL, *q = NULL;
		int tokenCount = 0;
		BOOL foundValue = FALSE;

		if (strlen(line) == 0) continue;
		else if (strcmp("\n", line) == 0 || strcmp("\r\n", line) == 0) continue;
		else p = line;

		// Read a token at a time.
		while (*p) {
			q = token;

			while((*p != ' ' && *p != '\t' && *p != ',' && *p != '\r' && *p != '\n') && *p) {
				*q = *p;
				q++; p++;
			}
			if (*p) p++; 
			*q = '\0';

			if (strlen(token) > 0) {
				tokenCount++;

				// Found the column that is wanted
				if (tokenCount == ops->mpheno + 2) { // First 2 columns ignored.
					foundValue = TRUE;
					break;
				}
			}
		}

		if (!foundValue) {
			if (fin != NULL) fclose(fin);
			sprintf(msg, "Unable to find specified column index (i=%i) in the PHE file (%s).", ops->mpheno, ops->szPheFileName);
			if(ops->verbose) error(msg);
			return 0;
		}

		// Are we using 0/1 coding
		if (ops->coding01) {
			if (strcmp(token, "1") == 0) strcpy(token, "2");
			else if (strcmp(token, "0") == 0) strcpy(token, "1");
			else strcpy(token, "0");
		}

		if (sampleIndex < ped->nSample) {
			ped->sample[sampleIndex]->phenotype = atof(token);
			ped->sample[sampleIndex]->missing = FALSE;

			// Check if 'qt' trait.
			if (!ops->qt) {
				if (strcmp(token, ops->missing_phenotype) != 0 && strcmp(token, "0") != 0 && strcmp(token, "1") != 0 && strcmp(token, "2") != 0) 
				{
					ops->qt = TRUE;
					ops->bt = FALSE;
				}
			}
		} else {
			if (fin != NULL) fclose(fin);
			if(ops->verbose) error("Lines in the PHE file do not match number of lines in the PED file.");
			return 0;
		}
		sampleIndex++;
	}
	if (fin != NULL) fclose(fin);

	// If a binary trait, now make 0 missing also.
	if (ops->bt) {
		for (i = 0; i < ped->nSample; i++) {
			if (ped->sample[i]->phenotype == missing_phenotype) ped->sample[i]->missing = TRUE;
		}

		// As binary trait, also apply the affected status to the phenotype data.
		phefile_aff_coding(ped);
	}

	return 1;
}

int phefile_aff_coding(struct pedfile *ped) {
	int i;

	if (ped == NULL) return 0;

	for (i = 0; i < ped->nSample; i++) {
		if (ped->sample[i]->phenotype == 2 && !ped->sample[i]->missing) 
			ped->sample[i]->aff = TRUE;
		else
			ped->sample[i]->aff = FALSE;
	}

	return 1;
}

double* phefile_load(struct selected_options *ops, int nSamples, BOOL *phe_missing) {
	char line[MAX_LINE_LENGTH], token[MAX_TOKEN_LENGTH];
	double *values = NULL;
	double missing_phenotype = 0;
	int i = 0, sampleIndex = 0;
	BOOL jobDone = FALSE;

	if (ops == NULL || nSamples <= 0 || phe_missing == NULL) return NULL;
	else if (ops->mpheno < 1 || strlen(ops->szPheFileName) == 0) return NULL; 

	missing_phenotype = atof(ops->missing_phenotype);
	init_bool_array(phe_missing, nSamples, FALSE);
	
	// Assume binary trait unless we find out otherwise.
	ops->bt = TRUE;
	ops->qt = FALSE;

	// Set all initial phenotypes to missing until read from the text file.
	values = create_vector(nSamples, missing_phenotype); 
	if (values != NULL) {
		FILE *in = fopen(ops->szPheFileName, "r");
		if (in != NULL) {
			// Find the column name if specified
			if (strlen(ops->szPheColumnName) > 0) {
				if(fgets(line,MAX_LINE_LENGTH, in)) {
					BOOL colFound = FALSE;
					int nTokens = get_number_of_tokens(line, ' '), i;
					for (i = 0; i < nTokens; i++) {
						get_token(line, token, i, ' ');
						if (strcmp(token, ops->szPheColumnName) == 0) {
							colFound = TRUE;
							ops->mpheno = i - 2 + 1;  // First 2 columns of FID and IID ignored, PHENO indices start at 1.
						}
					}
					if (!colFound) {
						free(values);
						init_bool_array(phe_missing, nSamples, FALSE);
						return NULL;
					}
				} 
			}
			while(fgets(line,MAX_LINE_LENGTH, in) && sampleIndex < nSamples) {
				char *p = NULL, *q = NULL;
				int tokenCount = 0;
				BOOL foundValue = FALSE;

				if (strlen(line) == 0) continue;
				else if (strcmp("\n", line) == 0 || strcmp("\r\n", line) == 0) continue;
				else p = line;

				// Read a token at a time.
				while (*p) {
					q = token;

					while((*p != ' ' && *p != '\t' && *p != ',' && *p != '\r' && *p != '\n') && *p) {
						*q = *p;
						q++; p++;
					}
					if (*p) p++; 
					*q = '\0';

					if (strlen(token) > 0) {
						tokenCount++;

						// Found the column that is wanted
						if (tokenCount == ops->mpheno + 2) { // First 2 columns ignored.
							foundValue = TRUE;
							break;
						}
					}
				}

				if (!foundValue) break;
				else {
					// Are we using 0/1 coding
					if (ops->coding01) {
						if (strcmp(token, "1") == 0) strcpy(token, "2");
						else if (strcmp(token, "0") == 0) strcpy(token, "1");
						else strcpy(token, "0");
					}
				}

				// Set the PHE values
				values[sampleIndex] = atof(token);

				// Check if 'qt' trait.
				if (!ops->qt) {
					if (strcmp(token, ops->missing_phenotype) != 0 && strcmp(token, "0") != 0 && strcmp(token, "1") != 0 && strcmp(token, "2") != 0) {
						ops->qt = TRUE;
						ops->bt = FALSE;
					}
				}

				sampleIndex++;
			}
			fclose(in);

			if (sampleIndex == nSamples) {
				jobDone = TRUE;

				// Apply the Affected status if binary trait to the 
				if (ops->bt) {
					for (i = 0; i < nSamples; i++) if (values[i] == missing_phenotype) phe_missing[i] = TRUE;
				}
			}
		}
	}

	if (!jobDone) {
		if (values != NULL) {
			free(values);
			values = NULL;
		}
	}
	
	return values;
}
